{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc6caafa",
   "metadata": {
    "id": "cc6caafa"
   },
   "source": [
    "# NVIDIA NIMs\n",
    "\n",
    "The `langchain-nvidia-ai-endpoints` package contains LangChain integrations building applications with models on \n",
    "NVIDIA NIM inference microservice. NIM supports models across domains like chat, embedding, and re-ranking models \n",
    "from the community as well as NVIDIA. These models are optimized by NVIDIA to deliver the best performance on NVIDIA \n",
    "accelerated infrastructure and deployed as a NIM, an easy-to-use, prebuilt containers that deploy anywhere using a single \n",
    "command on NVIDIA accelerated infrastructure.\n",
    "\n",
    "NVIDIA hosted deployments of NIMs are available to test on the [NVIDIA API catalog](https://build.nvidia.com/). After testing, \n",
    "NIMs can be exported from NVIDIA’s API catalog using the NVIDIA AI Enterprise license and run on-premises or in the cloud, \n",
    "giving enterprises ownership and full control of their IP and AI application.\n",
    "\n",
    "NIMs are packaged as container images on a per model basis and are distributed as NGC container images through the NVIDIA NGC Catalog. \n",
    "At their core, NIMs provide easy, consistent, and familiar APIs for running inference on an AI model.\n",
    "\n",
    "This example goes over how to use LangChain to interact with NVIDIA supported via the `ChatNVIDIA` class.\n",
    "\n",
    "For more information on accessing the chat models through this api, check out the [ChatNVIDIA](https://python.langchain.com/docs/integrations/chat/nvidia_ai_endpoints/) documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2be90a9",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e13eb331",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet langchain-nvidia-ai-endpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccff689e",
   "metadata": {
    "id": "ccff689e"
   },
   "source": [
    "## Setup\n",
    "\n",
    "**To get started:**\n",
    "\n",
    "1. Create a free account with [NVIDIA](https://build.nvidia.com/), which hosts NVIDIA AI Foundation models.\n",
    "\n",
    "2. Click on your model of choice.\n",
    "\n",
    "3. Under `Input` select the `Python` tab, and click `Get API Key`. Then click `Generate Key`.\n",
    "\n",
    "4. Copy and save the generated key as `NVIDIA_API_KEY`. From there, you should have access to the endpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "686c4d2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af0ce26b",
   "metadata": {},
   "source": [
    "## Working with NVIDIA API Catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "Jdl2NUfMhi4J",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jdl2NUfMhi4J",
    "outputId": "e9c4cc72-8db6-414b-d8e9-95de93fc5db4"
   },
   "outputs": [],
   "source": [
    "## Core LC Chat Interface\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "llm = ChatNVIDIA(model=\"mistralai/mixtral-8x7b-instruct-v0.1\")\n",
    "result = llm.invoke(\"Write a ballad about LangChain.\")\n",
    "print(result.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d35686b",
   "metadata": {},
   "source": [
    "## Working with NVIDIA NIMs\n",
    "When ready to deploy, you can self-host models with NVIDIA NIM—which is included with the NVIDIA AI Enterprise software license—and run them anywhere, giving you ownership of your customizations and full control of your intellectual property (IP) and AI applications.\n",
    "\n",
    "[Learn more about NIMs](https://developer.nvidia.com/blog/nvidia-nim-offers-optimized-inference-microservices-for-deploying-ai-models-at-scale/)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49838930",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "# connect to an embedding NIM running at localhost:8000, specifying a specific model\n",
    "llm = ChatNVIDIA(base_url=\"http://localhost:8000/v1\", model=\"meta/llama3-8b-instruct\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71d37987-d568-4a73-9d2a-8bd86323f8bf",
   "metadata": {},
   "source": [
    "## Stream, Batch, and Async\n",
    "\n",
    "These models natively support streaming, and as is the case with all LangChain LLMs they expose a batch method to handle concurrent requests, as well as async methods for invoke, stream, and batch. Below are a few examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fa5095-be72-47b0-8247-e9fac799435d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(llm.batch([\"What's 2*3?\", \"What's 2*6?\"]))\n",
    "# Or via the async API\n",
    "# await llm.abatch([\"What's 2*3?\", \"What's 2*6?\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75189ac6-e13f-414f-9064-075c77d6e754",
   "metadata": {},
   "outputs": [],
   "source": [
    "for chunk in llm.stream(\"How far can a seagull fly in one day?\"):\n",
    "    # Show the token separations\n",
    "    print(chunk.content, end=\"|\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a9a4122-7a10-40c0-a979-82a769ce7f6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "async for chunk in llm.astream(\n",
    "    \"How long does it take for monarch butterflies to migrate?\"\n",
    "):\n",
    "    print(chunk.content, end=\"|\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6RrXHC_XqWc1",
   "metadata": {
    "id": "6RrXHC_XqWc1"
   },
   "source": [
    "## Supported models\n",
    "\n",
    "Querying `available_models` will still give you all of the other models offered by your API credentials.\n",
    "\n",
    "The `playground_` prefix is optional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b8a312d-38e9-4528-843e-59451bdadbac",
   "metadata": {},
   "outputs": [],
   "source": [
    "ChatNVIDIA.get_available_models()\n",
    "# llm.get_available_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a407c6-e38b-4cfc-9a33-bcafadc18cf2",
   "metadata": {},
   "source": [
    "## Model types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "WMW79Iegqj4e",
   "metadata": {
    "id": "WMW79Iegqj4e"
   },
   "source": [
    "All of these models above are supported and can be accessed via `ChatNVIDIA`. \n",
    "\n",
    "Some model types support unique prompting techniques and chat messages. We will review a few important ones below.\n",
    "\n",
    "**To find out more about a specific model, please navigate to the API section of an AI Foundation model [as linked here](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/codellama-13b/api).**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d65053-59fe-40cf-a2d0-55d3dbb13585",
   "metadata": {},
   "source": [
    "### General Chat\n",
    "\n",
    "Models such as `meta/llama3-8b-instruct` and `mistralai/mixtral-8x22b-instruct-v0.1` are good all-around models that you can use for with any LangChain chat messages. Example below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5f7aee8-e90c-4d5a-ac97-0dd3d45c3f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful AI assistant named Fred.\"), (\"user\", \"{input}\")]\n",
    ")\n",
    "chain = prompt | ChatNVIDIA(model=\"meta/llama3-8b-instruct\") | StrOutputParser()\n",
    "\n",
    "for txt in chain.stream({\"input\": \"What's your name?\"}):\n",
    "    print(txt, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04146118-281b-42ef-b781-2fadeeeea6c8",
   "metadata": {},
   "source": [
    "### Code Generation\n",
    "\n",
    "These models accept the same arguments and input structure as regular chat models, but they tend to perform better on code-genreation and structured code tasks. An example of this is `meta/codellama-70b`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49aa569b-5f33-47b3-9edc-df58313eb038",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are an expert coding AI. Respond only in valid python; no narration whatsoever.\",\n",
    "        ),\n",
    "        (\"user\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "chain = prompt | ChatNVIDIA(model=\"meta/codellama-70b\") | StrOutputParser()\n",
    "\n",
    "for txt in chain.stream({\"input\": \"How do I solve this fizz buzz problem?\"}):\n",
    "    print(txt, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f465ff6-5922-41d8-8abb-1d1e4095cc27",
   "metadata": {},
   "source": [
    "## Multimodal\n",
    "\n",
    "NVIDIA also supports multimodal inputs, meaning you can provide both images and text for the model to reason over. An example model supporting multimodal inputs is `nvidia/neva-22b`.\n",
    "\n",
    "\n",
    "These models accept LangChain's standard image formats, and accept `labels`, similar to the Steering LLMs above. In addition to `creativity`, `complexity`, and `verbosity`, these models support a `quality` toggle.\n",
    "\n",
    "Below is an example use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26625437-1695-440f-b792-b85e6add9a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import requests\n",
    "\n",
    "image_url = \"https://www.nvidia.com/content/dam/en-zz/Solutions/research/ai-playground/nvidia-picasso-3c33-p@2x.jpg\"  ## Large Image\n",
    "image_content = requests.get(image_url).content\n",
    "\n",
    "IPython.display.Image(image_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbbe57c-27a5-4cbb-b967-19c4e7d29fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "llm = ChatNVIDIA(model=\"nvidia/neva-22b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ddcb8f1-9cd8-4376-963d-af61c29b2a3c",
   "metadata": {},
   "source": [
    "#### Passing an image as a URL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "432ea2a2-4d39-43f8-a236-041294171f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "llm.invoke(\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=[\n",
    "                {\"type\": \"text\", \"text\": \"Describe this image:\"},\n",
    "                {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n",
    "            ]\n",
    "        )\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0573dd1f-9a17-4c99-ab2a-8d930b89d283",
   "metadata": {},
   "source": [
    "#### Passing an image as a base64 encoded string"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be1688a5",
   "metadata": {},
   "source": [
    "At the moment, some extra processing happens client-side to support larger images like the one above. But for smaller images (and to better illustrate the process going on under the hood), we can directly pass in the image as shown below: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c58f1dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import requests\n",
    "\n",
    "image_url = \"https://picsum.photos/seed/kitten/300/200\"\n",
    "image_content = requests.get(image_url).content\n",
    "\n",
    "IPython.display.Image(image_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c721629-42eb-4006-bf68-0296f7925ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "\n",
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "## Works for simpler images. For larger images, see actual implementation\n",
    "b64_string = base64.b64encode(image_content).decode(\"utf-8\")\n",
    "\n",
    "llm.invoke(\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=[\n",
    "                {\"type\": \"text\", \"text\": \"Describe this image:\"},\n",
    "                {\n",
    "                    \"type\": \"image_url\",\n",
    "                    \"image_url\": {\"url\": f\"data:image/png;base64,{b64_string}\"},\n",
    "                },\n",
    "            ]\n",
    "        )\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba958424-28d7-4bc2-9c8e-bd571066853f",
   "metadata": {},
   "source": [
    "#### Directly within the string\n",
    "\n",
    "The NVIDIA API uniquely accepts images as base64 images inlined within `<img/>` HTML tags. While this isn't interoperable with other LLMs, you can directly prompt the model accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00c06a9a-497b-4192-a842-b075e27401aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "base64_with_mime_type = f\"data:image/png;base64,{b64_string}\"\n",
    "llm.invoke(f'What\\'s in this image?\\n<img src=\"{base64_with_mime_type}\" />')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e61d868",
   "metadata": {},
   "source": [
    "#### **Advanced Use Case:** Forcing Payload \n",
    "\n",
    "You may notice that some newer models may have strong parameter expectations that the LangChain connector may not support by default. For example, we cannot invoke the [Kosmos](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/kosmos-2) model at the time of this notebook's latest release due to the lack of a streaming argument on the server side: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d143e0d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "kosmos = ChatNVIDIA(model=\"microsoft/kosmos-2\")\n",
    "\n",
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "# kosmos.invoke(\n",
    "#     [\n",
    "#         HumanMessage(\n",
    "#             content=[\n",
    "#                 {\"type\": \"text\", \"text\": \"Describe this image:\"},\n",
    "#                 {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n",
    "#             ]\n",
    "#         )\n",
    "#     ]\n",
    "# )\n",
    "\n",
    "# Exception: [422] Unprocessable Entity\n",
    "# body -> stream\n",
    "#   Extra inputs are not permitted (type=extra_forbidden)\n",
    "# RequestID: 35538c9a-4b45-4616-8b75-7ef816fccf38"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e230b70",
   "metadata": {},
   "source": [
    "For a simple use case like this, we can actually try to force the payload argument of our underlying client by specifying the `payload_fn` function as follows: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0925b2b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def drop_streaming_key(d):\n",
    "    \"\"\"Takes in payload dictionary, outputs new payload dictionary\"\"\"\n",
    "    if \"stream\" in d:\n",
    "        d.pop(\"stream\")\n",
    "    return d\n",
    "\n",
    "\n",
    "## Override the payload passthrough. Default is to pass through the payload as is.\n",
    "kosmos = ChatNVIDIA(model=\"microsoft/kosmos-2\")\n",
    "kosmos.client.payload_fn = drop_streaming_key\n",
    "\n",
    "kosmos.invoke(\n",
    "    [\n",
    "        HumanMessage(\n",
    "            content=[\n",
    "                {\"type\": \"text\", \"text\": \"Describe this image:\"},\n",
    "                {\"type\": \"image_url\", \"image_url\": {\"url\": image_url}},\n",
    "            ]\n",
    "        )\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6e1758",
   "metadata": {},
   "source": [
    "For more advanced or custom use-cases (i.e. supporting the diffusion models), you may be interested in leveraging the `NVEModel` client as a requests backbone. The `NVIDIAEmbeddings` class is a good source of inspiration for this. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "137662a6",
   "metadata": {
    "id": "137662a6"
   },
   "source": [
    "## Example usage within a Conversation Chains"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79efa62d",
   "metadata": {
    "id": "79efa62d"
   },
   "source": [
    "Like any other integration, ChatNVIDIA is fine to support chat utilities like conversation buffers by default. Below, we show the [LangChain ConversationBufferMemory](https://python.langchain.com/docs/modules/memory/types/buffer) example applied to the `mistralai/mixtral-8x22b-instruct-v0.1` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "082ccb21-91e1-4e71-a9ba-4bff1e89f105",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd2c6bc1",
   "metadata": {
    "id": "fd2c6bc1"
   },
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationChain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "\n",
    "chat = ChatNVIDIA(\n",
    "    model=\"mistralai/mixtral-8x22b-instruct-v0.1\",\n",
    "    temperature=0.1,\n",
    "    max_tokens=100,\n",
    "    top_p=1.0,\n",
    ")\n",
    "\n",
    "conversation = ConversationChain(llm=chat, memory=ConversationBufferMemory())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f644ff28",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 268
    },
    "id": "f644ff28",
    "outputId": "bae354cc-2118-4e01-ce20-a717ac94d27d"
   },
   "outputs": [],
   "source": [
    "conversation.invoke(\"Hi there!\")[\"response\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "uHIMZxVSVNBC",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 284
    },
    "id": "uHIMZxVSVNBC",
    "outputId": "79acc89d-a820-4f2c-bac2-afe99da95580"
   },
   "outputs": [],
   "source": [
    "conversation.invoke(\"I'm doing well! Just having a conversation with an AI.\")[\n",
    "    \"response\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "LyD1xVKmVSs4",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 350
    },
    "id": "LyD1xVKmVSs4",
    "outputId": "a1714513-a8fd-4d14-f974-233e39d5c4f5"
   },
   "outputs": [],
   "source": [
    "conversation.invoke(\"Tell me about yourself.\")[\"response\"]"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
